feat: Wire BatchWriteRecord and ListRecords into ingest_dataframe#6026
feat: Wire BatchWriteRecord and ListRecords into ingest_dataframe#6026chiragvp-aws wants to merge 1 commit into
Conversation
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
| for index in range(1, len(row)): | ||
| feature_name = data_frame.columns[index - 1] |
There was a problem hiding this comment.
Not sure I'm following the 1-based indexing used here
There was a problem hiding this comment.
In dataframe, column 0 is the index of the row which should be discarded. We have similar logic for _ingest_data using putRecord.
| return record | ||
|
|
||
| @staticmethod | ||
| def _ingest_batch_write( |
There was a problem hiding this comment.
Lots of chances to accidentally author off-by-1 errors as we're doing a lot of indexing/slicing arrays here. Do we have sufficient unit tests to cover a wide range of scenarios?
There was a problem hiding this comment.
We have good number of unit tests. I will check if we have some missing scenarios here.
| if response.unprocessed_entries: | ||
| for i, entry in enumerate(response.unprocessed_entries): | ||
| # unprocessed_entries are BatchWriteRecordEntry objects; | ||
| # find their position in the original entries list |
There was a problem hiding this comment.
Should we find their position in the original entries list, or in the original data frame?
There was a problem hiding this comment.
Not required as we are building entries for request body at line 451.
| idx = entries.index(entry) | ||
| failed_rows.append(row_indices[idx]) | ||
| except ValueError: | ||
| # If we can't find the exact entry, mark by position |
There was a problem hiding this comment.
When would we run into this except block?
| @@ -0,0 +1,362 @@ | |||
| # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. | |||
| # Licensed under the Apache License, Version 2.0 | |||
| """Unit tests for BatchWriteRecord and ListRecords wiring.""" | |||
There was a problem hiding this comment.
Why couple them into the same test module? Should we separate them into 2 modules?
| assert mgr.failed_rows == [] | ||
|
|
||
| @patch("sagemaker.mlops.feature_store.ingestion_manager_pandas.CoreFeatureGroup") | ||
| def test_batch_write_partial_failure_maps_to_row(self, mock_fg_class, feature_definitions): |
There was a problem hiding this comment.
Are there any other kinds of partial failures we need to test? Perhaps a scenario with only unprocessed_entries, one with only errors, one with both?
There was a problem hiding this comment.
I'll see if we can add more tests.
d0b514f to
9eda65d
Compare
GitHub Actions
Inactive
GitHub Actions
Inactive
9eda65d to
f79f93a
Compare
GitHub Actions
Inactive
GitHub Actions
Inactive
f79f93a to
5f513fb
Compare
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
|
All 16 errors (both new files) fail at setup with: Please follow that file's pattern: generate a unique name with unique_name_from_base("integ-test-fg"), FeatureGroup.create(...) → fg.wait_for_status("Created") → (for list_records, ingest data) → run assertions → tear down in a finally via FeatureGroup.get(...).delete(). Reuse the shared role/bucket/region fixtures.(test_list_records_non_existent can stay as-is — it's meant to test the not-found path.) |
02b3971 to
b2afe55
Compare
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
- Add use_batch_write_record=False flag to ingest_dataframe() (Proposal C) - Implement _ingest_batch_write() with 25 records per API call - Map partial failures (response.errors) back to specific row indices - Add list_records() utility function with pagination support - Export list_records from feature_store __init__.py Tested: 50 integration tests + 17 unit tests
b2afe55 to
64f582b
Compare
GitHub Actions
Inactive
GitHub Actions
Inactive
GitHub Actions
Inactive
3 participants
Tested: 50 integration tests + 17 unit tests
Issue #, if available:
Description of changes:
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.